from torch import nn
import math
import torch.utils.model_zoo as model_zoo
import torch
import torch.nn.functional as F
from copy import deepcopy

from functools import partial
def initialize_weights(*models):
    for model in models:
        for module in model.modules():
            if isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear):
                nn.init.kaiming_normal(module.weight)
                if module.bias is not None:
                    module.bias.data.zero_()
            elif isinstance(module, nn.BatchNorm2d):
                module.weight.data.fill_(1)
                module.bias.data.zero_()
def conv3x3(in_planes, out_planes, stride=1):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    squeeze = 16

    def __init__(self, inplanes, outplanes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, outplanes//self.squeeze, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(outplanes//self.squeeze)
        self.conv2 = nn.Conv2d(outplanes//self.squeeze, outplanes//self.squeeze, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(outplanes//self.squeeze)
        self.conv3 = nn.Conv2d(outplanes//self.squeeze, outplanes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(outplanes)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out

class Bottleneck_up(nn.Module):
    squeeze = 16

    def __init__(self, inplanes, outplanes, stride=1, upsample=None):
        super(Bottleneck_up, self).__init__()
        self.conv1 = nn.Conv2d(inplanes, inplanes//self.squeeze, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(inplanes//self.squeeze)
        self.conv2 = nn.ConvTranspose2d(inplanes//self.squeeze, inplanes//self.squeeze, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(inplanes//self.squeeze)
        self.conv3 = nn.Conv2d(inplanes//self.squeeze, outplanes, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(outplanes)
        self.relu = nn.ReLU(inplace=True)
        self.upsample = upsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.upsample is not None:
            residual = self.upsample(x)

        out += residual
        out = self.relu(out)

        return out


class _EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, num_block, stride):
        super(_EncoderBlock, self).__init__()
        self.inplanes = in_channels
        self.layer1 = self._make_layer(Bottleneck, out_channels, num_block, stride)

    def forward(self, x):
        return self.layer1(x)

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)


class _DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels, num_block, stride):
        super(_DecoderBlock, self).__init__()
        self.inplanes = in_channels
        self.layer1 = self._make_layer(Bottleneck_up, out_channels, num_block,stride)

    def forward(self, x):
        return self.layer1(x)

    def _make_layer(self, block, planes, blocks, stride=1):

        upsample = nn.Sequential(
            nn.ConvTranspose2d(self.inplanes, planes, kernel_size=1, stride=stride, bias=False),
            nn.BatchNorm2d(planes),
        )

        layers = []

        for i in range(1, blocks):
            layers.append(block(self.inplanes, self.inplanes))
        layers.append(block(self.inplanes, planes, stride, upsample))

        return nn.Sequential(*layers)

class fuckSeg(nn.Module):
    '''
    0.830
    '''
    def __init__(self, num_classes):
        super(fuckSeg,self).__init__()
        self.enc0 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                      bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.enc1 = _EncoderBlock(64, 128,num_block=3, stride=1)  #1/2
        self.enc2 = _EncoderBlock(128, 256, num_block=4, stride=2) # 1/4
        self.enc3 = _EncoderBlock(256, 512,num_block=6, stride=2)  # 1/8
        self.enc4 = _EncoderBlock(512, 1024,num_block=3, stride=2)  # 1/16

        self.dec1 = _DecoderBlock(1024, 512,num_block=3, stride=2)  # 1/8  en4
        self.dec2 = _DecoderBlock(512, 256,num_block=4, stride=2)  # 1/4  de1 cat en3
        self.dec3 = _DecoderBlock(256, 128,num_block=3, stride=2)  # 1/2  de2 cat en2
        self.dec4 = _DecoderBlock(128, num_classes, num_block=3, stride=2)  # 1 de3 cat en1
        initialize_weights(self)
    def forward(self, x):
        x = self.enc0(x)  # 1/2
        x = self.enc1(x)  # 1/2
        x = self.enc2(x)  # 1/4
        x = self.enc3(x)  # 1/8
        x = self.enc4(x)  # 1/16

        x = self.dec1(x) # 1/8
        x = self.dec2(x) # 1/4
        x = self.dec3(x) # # 1/2
        x = self.dec4(x) # # 1
        x = x.squeeze(1)
        return x

class fuckSeg2(nn.Module):
    def __init__(self, num_classes):
        super(fuckSeg2,self).__init__()
        self.enc0 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                      bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.enc1 = _EncoderBlock(64, 128,num_block=3, stride=1)  #1/2
        self.enc2 = _EncoderBlock(128, 256, num_block=4, stride=2) # 1/4
        self.enc3 = _EncoderBlock(256, 512,num_block=6, stride=2)  # 1/8
        self.enc4 = _EncoderBlock(512, 1024,num_block=3, stride=2)  # 1/16

        self.dec1 = _DecoderBlock(1024, 512,num_block=3, stride=2)  # 1/8  en4
        self.dec2 = _DecoderBlock(512+512, 512,num_block=4, stride=2)  # 1/4  de1 cat en3
        self.dec3 = _DecoderBlock(512+256, 256,num_block=3, stride=2)  # 1/2  de2 cat en2
        self.dec4 = _DecoderBlock(256+128, num_classes, num_block=3, stride=2)  # 1 de3 cat en1
        initialize_weights(self)
    def forward(self, x):
        x = self.enc0(x)  # 1/2
        en1 = self.enc1(x)  # 1/2
        en2 = self.enc2(en1)  # 1/4
        en3 = self.enc3(en2)  # 1/8
        en4 = self.enc4(en3)  # 1/16

        x = self.dec1(en4) # 1/8
        x = self.dec2(torch.cat([en3, F.upsample(x, en3.size()[2:], mode='bilinear')],dim=1)) # 1/4
        x = self.dec3(torch.cat([en2, F.upsample(x, en2.size()[2:], mode='bilinear')],dim=1)) # # 1/2
        x = self.dec4(torch.cat([en1, F.upsample(x, en1.size()[2:], mode='bilinear')],dim=1)) # # 1
        x = x.squeeze(1)
        return x



# class UNet(nn.Module):
#     def __init__(self, num_classes):
#         super(UNet, self).__init__()
#         self.enc1 = _EncoderBlock(3, 64)

# model = _EncoderBlock(in_channels=3, mid_channels=16, num_block=3)
# print model

if __name__=='__main__':
    import numpy as np
    # model = _DecoderBlock(in_channels=512, out_channels=256, num_block=2,stride=1)
    # model = _EncoderBlock(in_channels=256, out_channels=512, num_block=2, stride=1)
    model=fuckSeg2(num_classes=2)
    model.load_state_dict(torch.load('/home/hszc/zhangchi/models/FigureSeg/fuckSeg/weights-73-84-[0.840].pth'))
    for item in model.state_dict():
        print model.state_dict()[item]
        # break

    # torch.save(model.state_dict(),'./testnew.pth')

    # print model

    x = torch.FloatTensor(1,3,225,225)
    # x = torch.from_numpy(np.arange(4).reshape(1,1,2,2)).float()
    # print x
    x = torch.autograd.Variable(x)
    y = model(x)
    # #
    print y.size()